#!/usr/bin/env python

import numpy as np
from scipy.stats import dirichlet, rv_discrete, bernoulli

from information import *

alpha = np.ones(10)
pi = dirichlet(alpha).rvs(size=1)[0]

px = rv_discrete(name='PX', values=(np.arange(10), pi))

z = bernoulli.rvs(p=0.5, size=100)

x = px.rvs(size=100)


def ml(z):

    x0 = x[z==0]
    x1 = x[z==1]

    H0 = entropy(x0)

    H1= entropy(x1)

    N0, N1 = len(x0), len(x1)
    N = N0+N1
    Hz = entropy(z) * N
    return H0 * N0 + H1 * N1 + Hz

lmax = 0
for _ in range(100):
    l = ml(z)
    if l > lmax:
        lmax = l
        print(_, lmax)
        zmax = z
for xi, zi in zip(x,z):
    print(xi, zi)

from collections import Counter
print(Counter(x[z==0]), Counter(x[z==1]))

